import os

from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data import SubsetRandomSampler
from torchvision.datasets import ImageFolder
from preprocess.transforms import ToTensor,Compose,RandomRotate,Center_Crop,RandomAffine,ColorJitter
from torchvision import transforms

# https://pytorch.org/vision/master/_modules/torchvision/transforms/functional.html
class TrainDataset(Dataset):
    def __init__(self, image_path_list):
        # generate image path list

        self.image_path_list = image_path_list

        self.transforms =  Compose([

            ToTensor(),
            RandomRotate(30),
            RandomAffine(degrees=10, translate=(0.1,0.1)),
            Center_Crop(160),
            # ColorJitter(brightness=0.1,contrast=0.1,saturation=0.1)
        ])
        self .img_transform = transforms.Compose(
            [
                # transforms.RandomErasing(scale=(0.02, 0.33)),
            ]
        )

    def __getitem__(self, index):
        image_path,mask_path = self.image_path_list[index]
        image = Image.open(image_path)
        mask = Image.open(mask_path)

        # get label from dirname
        label = int(image_path.split("/")[-3])
        label = torch.tensor(label)

        image,mask = self.transforms(image,mask)
        image = self.img_transform(image)
        return image,mask,label

    def __len__(self):
        return len(self.image_path_list)
